#Copyright (c) 2020 Ocado. All Rights Reserved.

import sys, os, pygame, argparse
from PIL import Image
import numpy as np
import cv2

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)))
# print(sys.path)
# print(os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir)))

from amrrt.space import StateSpace
from amrrt.diffusion_map import DiffusionMap, GridGraph
from amrrt.metrics import EuclideanMetric, DiffusionMetric, GeodesicMetric
from amrrt.planners2 import AMRRTPlanner, RTRRTPlanner,SYRTRRTPlanner
from amrrt.approx_work2 import approx

img_path='./resources/maze.png'
image_size = 800
#ratio = 0.015#
ratio = 0.02#(maze)
#图大了把ratio调小
Obs_radius=2#maze
#Obs_radius=5

Obs_speed=3


def display(screen, planner, agent_pos, path, image, scale,planner_type):
    screen.blit(image, (0, 0))
    for u in planner.tree.edges:#u是一个节点
        for v in planner.tree.edges[u]:#edges[u]也就是v,是u的孩子
            a = (int(u.pos[0]*scale), int(u.pos[1]*scale))
            b = (int(v.pos[0]*scale), int(v.pos[1]*scale))
            pygame.draw.line(screen, (0,0,0), a, b, 2)# 孩子和父亲连起来
    display_path = [planner.space.create_state(agent_pos)] + path
    for i in range(len(display_path)-1):
        a = (int(display_path[i].pos[0]*scale), int(display_path[i].pos[1]*scale))
        b = (int(display_path[i+1].pos[0]*scale), int(display_path[i+1].pos[1]*scale))
        pygame.draw.line(screen, (30,30,255), a, b, 8)# 画路径
    for obstacle in planner.space.dynamic_obstacles:
        pygame.draw.circle(screen, (255,0,0), tuple((obstacle.pos*scale).astype(int)), int(obstacle.radius*scale))
        # 代码遍历每一个动态障碍物，并使用红色圆圈将它们显示在窗口上。
    for root in planner.root_queue:
        pygame.draw.circle(screen, (30,255,30), tuple((root.pos*scale).astype(int)), 4)
        # 显示从根节点重布线的所有节点
    for rand in planner.rand_queue:
        pygame.draw.circle(screen, (255,30,30), tuple((rand.pos * scale).astype(int)), 4)
        # 显示从随机节点重布线的所有节点
    pygame.draw.circle(screen, (255,30,30), tuple((agent_pos*scale).astype(int)), 4)
    # 代码绘制机器人的当前位置，使用红色圆圈表示。
    if planner.goal is not None:
        pygame.draw.circle(screen, (30,255,30), tuple((planner.goal.pos*scale).astype(int)), 10)
        # 如果规划器对象 planner 中存储了目标位置信息，则代码将目标位置用绿色圆圈表示。

    # 画剖分线
    if planner_type=="syrtrrt" :
        for cut_line in planner.Approx.cutlinelist:
            pygame.draw.line(screen, (0,255,0), planner.Approx.point_map[cut_line[0]],planner.Approx.point_map[cut_line[1]], 1)

    pygame.display.update()


def convex(path):
    robotsize = 0.6
    thresh = 128

    IMGmap = cv2.imread(path)

    # #调整图像大小
    height, width, channels = IMGmap.shape
    scale = image_size/min(height,width)#比例，通用的大小为500
    new_height,new_width = int(height*scale), int(width*scale)#长宽中至少有一个是500
    IMGmap = cv2.resize(IMGmap, (new_width, new_height))  # 调整图形大小

    # cv2.imshow("1", IMGmap)
    Approx=approx()
    returndata = Approx.approx_work(IMGmap, ratio, robotsize, thresh)
    #print(Approx.approxs)
    # 输入处理过后的路径，比例，机器人大小，
    for i in range(0, len(returndata[0])):  # 遍历连通域,大多数情况下只有一个连通域,即len(returndata[0]）是1
        Approx.find_concave(returndata[0][i])
        #print(returndata[0][i])
        Approx.Startcut()
        Approx.GetPloygon()
        Approx.DrawBImap(False,IMGmap)

    return Approx,IMGmap

def Judgearea(planner,Approx,pos,scale,Obs_radius):
    planner.obstacles_area.add(Approx.IMGmap[int(pos[1] * scale), int(pos[0] * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int((pos[1] + Obs_radius) * scale), int(pos[0] * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int((pos[1] - Obs_radius) * scale), int(pos[0] * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int(pos[1] * scale), int((pos[0] + Obs_radius) * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int(pos[1] * scale), int((pos[0] - Obs_radius) * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int((pos[1] + Obs_radius) * scale), int((pos[0] + Obs_radius) * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int((pos[1] - Obs_radius) * scale), int((pos[0] + Obs_radius) * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int((pos[1] + Obs_radius) * scale), int((pos[0] - Obs_radius) * scale)][0])
    planner.obstacles_area.add(Approx.IMGmap[int((pos[1] - Obs_radius) * scale), int((pos[0] - Obs_radius) * scale)][0])


def visualiser(space, planner_type, metric_type):
    #space经过了StateSpace，主要生成一个二值化的矩阵
    pygame.init()

    # 剖分
    if planner_type == "syrtrrt":
        print("开始剖分")
        Approx,IMGmap=convex(img_path)
        print("剖分完毕")


    # grid_graph = GridGraph(space)#不知道为什么，这里要初始化巨久
    # 这玩意算高级辅助指标的时候用的到，单用rtrrt可以注释掉

    #因为导入了from amrrt.diffusion_map import DiffusionMap, GridGraph，可以直接用
    if metric_type == "euclidean" : assisting_metric = EuclideanMetric()    #辅助指标是EuclideanMetric()
    elif metric_type == "diffusion" : assisting_metric = DiffusionMetric(DiffusionMap(space, grid_graph=grid_graph))
    else : assisting_metric = GeodesicMetric(grid_graph)
    planner = None
    agent_pos = None
    image = space.display_image()   #调用了StateSpace->GridEnviroment.display_image,生成灰度图（RGB）
    scale = image_size/min(image.size[0],image.size[1])#比例，通用的大小为800
    width, height = int(image.size[0]*scale), int(image.size[1]*scale)#长宽中至少有一个是800
    screen = pygame.display.set_mode((width, height))#用pygame创建了一个窗口
    image = image.resize((width, height), resample=Image.NEAREST)#调整图形大小，resample 参数是可选的，指定了用于调整图像大小的算法。
    image = pygame.image.frombuffer(image.tobytes(), image.size, image.mode)

    #这样就可以渲染了
    #显示用的image但计算用的是space

    while planner is None:
        for event in pygame.event.get():
            #这里在等你输入起点
            if event.type == pygame.QUIT:
                sys.exit()
            if event.type == pygame.MOUSEBUTTONDOWN:
                pos = np.array(pygame.mouse.get_pos()) / scale#得到的点按比例转化，这里由 输出图像image 转化为 实际图像space
                if planner_type == "rtrrt" :
                    planner = RTRRTPlanner(space, pos, assisting_metric=assisting_metric)
                #初始化规划器，树啊什么的都在里面了
                elif planner_type == "syrtrrt":
                    planner = SYRTRRTPlanner(space, pos, Approx,scale,assisting_metric=assisting_metric)
                    planner.startconvex=Approx.IMGmap[int(pos[1]*scale),int(pos[0]*scale)][0]
                    print(Approx.IMGmap[int(pos[1]*scale),int(pos[0]*scale)])
                else :
                    planner = AMRRTPlanner(space, pos, assisting_metric=assisting_metric)
                agent_pos = pos
        screen.blit(image, (0, 0))

        #将一个图像（image）绘制到屏幕（screen）上，位置为左上角坐标为 (0, 0)。
        #blit() 方法会将一个 Surface 对象绘制到另一个 Surface 对象上，并返回一个表示被绘制部分的矩形。
        pygame.display.update()
        """
        pygame.display.update() 是 Pygame 中的一个函数，它的作用是更新屏幕显示。

        在 Pygame 中，当你通过绘制（如 screen.blit()）更新了游戏画面后，画面并不会立即显示在屏幕上，而是存储在缓冲区中。只有调用 pygame.display.update() 函数，才能将缓冲区中的内容刷新到屏幕上，使得玩家可以看到最新的游戏画面。

        注意，如果你只想更新部分区域的屏幕，而不是整个屏幕，你可以传递一个矩形参数给 update() 函数。例如，pygame.display.update(rect) 会更新 rect 矩形区域的屏幕显示。

        需要注意的是，在使用 update() 函数之前，需要保证屏幕对象已经被创建。一般来说，你需要先调用 pygame.init() 来初始化 Pygame，然后创建屏幕对象（screen），再调用 update() 函数来刷新屏幕。
        """
    stop=False
    Obs_num = 0
    old_goal=None
    while True:

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
            if pygame.key.get_pressed()[pygame.K_e]:
                stop = not stop
                # 强行停下来
            if pygame.key.get_pressed()[pygame.K_w]:
                dynamic_obs = planner.space.dynamic_obstacles.pop()
                dynamic_obs.pos[0] = dynamic_obs.pos[0]
                dynamic_obs.pos[1] = dynamic_obs.pos[1] - Obs_speed
                planner.space.dynamic_obstacles.append(dynamic_obs)
                if planner_type == "syrtrrt":
                    Judgearea(planner, Approx, dynamic_obs.pos, scale, Obs_radius)
                    print(planner.obstacles_area)
            if pygame.key.get_pressed()[pygame.K_s]:
                dynamic_obs = planner.space.dynamic_obstacles.pop()
                dynamic_obs.pos[0] = dynamic_obs.pos[0]
                dynamic_obs.pos[1] = dynamic_obs.pos[1] + Obs_speed
                planner.space.dynamic_obstacles.append(dynamic_obs)
                if planner_type == "syrtrrt":
                    Judgearea(planner, Approx, dynamic_obs.pos, scale, Obs_radius)
                    print(planner.obstacles_area)
            if pygame.key.get_pressed()[pygame.K_a]:
                dynamic_obs = planner.space.dynamic_obstacles.pop()
                dynamic_obs.pos[0] = dynamic_obs.pos[0] - Obs_speed
                dynamic_obs.pos[1] = dynamic_obs.pos[1]
                planner.space.dynamic_obstacles.append(dynamic_obs)
                if planner_type == "syrtrrt":
                    Judgearea(planner, Approx, dynamic_obs.pos, scale, Obs_radius)
                    print(planner.obstacles_area)
            if pygame.key.get_pressed()[pygame.K_d]:
                dynamic_obs=planner.space.dynamic_obstacles.pop()
                dynamic_obs.pos[0]=dynamic_obs.pos[0]+Obs_speed
                dynamic_obs.pos[1] = dynamic_obs.pos[1]
                planner.space.dynamic_obstacles.append(dynamic_obs)
                if planner_type == "syrtrrt":
                    Judgearea(planner, Approx, dynamic_obs.pos, scale, Obs_radius)
                    print(planner.obstacles_area)


            if event.type == pygame.MOUSEBUTTONDOWN:
                pos = np.array(pygame.mouse.get_pos()) / scale
                # print(Approx.IMGmap[int(pos[1] * scale), int(pos[0] * scale)])
                if pygame.key.get_pressed()[pygame.K_r]:#添加障碍物
                    planner.add_dynamic_obstacle(pos, Obs_radius,Obs_num)
                    Obs_num+=1
                    # for obs in planner.space.dynamic_obstacles:
                    #     print(obs.Obs_num)
                    #print("障碍物所占区域")
                    if planner_type=="syrtrrt":
                        Judgearea(planner,Approx,pos,scale,Obs_radius)
                        print(planner.obstacles_area)
                else:
                    planner.set_goal(pos)
                    if planner_type == "syrtrrt":
                        planner.obstacles_area.add(Approx.IMGmap[int(pos[1] * scale), int(pos[0] * scale)][0])
                    if old_goal==None:
                        old_goal = Approx.IMGmap[int(pos[1] * scale), int(pos[0] * scale)][0]
                    if old_goal !=Approx.IMGmap[int(pos[1] * scale), int(pos[0] * scale)][0] and old_goal!=None:
                        planner.obstacles_area.remove(old_goal)
                        old_goal = Approx.IMGmap[int(pos[1] * scale), int(pos[0] * scale)][0]
        path = []

        waypoint = planner.plan(agent_pos).pos.copy()
        #路径规划的主体，返回起点位置（准确的说是路径第一个节点）
        path = planner.goal_path()
        path = [planner.space.create_state(agent_pos)] + path
        for i in range(1, len(path)):
            if stop == True:
                agent_pos = agent_pos
                break
            if not planner.space.dynamically_free_path(path[i-1], path[i]) :
                # 如果路径上有障碍物，停下来
                agent_pos = agent_pos
                # 把当前起点加入到rrt的树中
                agent_node=planner.space.create_state(agent_pos)
                xnearest = planner.tree.nearest(agent_node)
                nearby = planner.find_nodes_near(agent_node)
                planner._add_node(agent_node, xnearest, nearby)
                planner.set_root(agent_node)
                break
            else:
                agent_pos += min(1, 0.7/np.linalg.norm(waypoint-agent_pos)) * (waypoint-agent_pos)

        '''
        具体来说，它首先计算机器人与目标之间的距离 np.linalg.norm(waypoint-agent_pos)，然后将其作为一个权重来调整机器人的移动方向。
        这个权重是 min(1, 0.7/np.linalg.norm(waypoint-agent_pos))，它的值在 [0, 1] 范围内，用于缩小机器人的移动步长，以避免机器人在靠近目标时过度加速，从而导致运动不稳定。

        然后，代码使用调整后的移动方向 (waypoint-agent_pos)，乘以这个权重，并将结果加到当前位置 agent_pos 上，以计算机器人的下一个位置。
        这个操作相当于将机器人从当前位置沿着调整后的方向移动一定的距离，距离的大小由权重控制。
        '''
        display(screen, planner, agent_pos, path, image, scale,planner_type)


if __name__ == "__main__":
    # parser = argparse.ArgumentParser(description='AM-RRT* & RT-RRT* graphical visualiser')
    # parser.add_argument('image', type=str, help='Filename of an image containing the environment')
    # parser.add_argument('planner', choices=['rtrrt', 'amrrt'], help='Name of planner (choices: "rtrrt", "amrrt")')
    # parser.add_argument('metric_type', choices=['euclidean', 'diffusion', 'geodesic'], help='Name of assisting metric (choices: "euclidean", "diffusion", "geodesic")')
    # args = parser.parse_args()
    # visualiser(StateSpace.from_image(args.image), args.planner, args.metric_type)
    visualiser(StateSpace.from_image(img_path), 'syrtrrt', 'euclidean')



